{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VwlxkU9cEGPR"
      },
      "source": [
        "Copyright 2020 The dnn-predict-accuracy Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "x08yky7rytbD"
      },
      "source": [
        "# README\n",
        "\n",
        "This notebook contains code for training predictors of DNN accuracy. \n",
        "\n",
        "Contents:\n",
        "\n",
        "(1) Loading the Small CNN Zoo dataset\n",
        "\n",
        "(2) Figure 2 of the paper\n",
        "\n",
        "(3) Examples of training Logit-Linear / GBM / DNN predictors\n",
        "\n",
        "(4) Transfer of predictors across CNN collections\n",
        "\n",
        "(5) Various visualizations of CNN collections\n",
        "\n",
        "Code dependencies:\n",
        "Light-GBM package\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aj14NoLVykBz"
      },
      "outputs": [],
      "source": [
        "from __future__ import division\n",
        "\n",
        "import time\n",
        "import os\n",
        "import json\n",
        "import sys\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "from matplotlib import colors\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "from scipy import stats\n",
        "from tensorflow import keras\n",
        "from tensorflow.io import gfile\n",
        "import lightgbm as lgb\n",
        "\n",
        "DATAFRAME_CONFIG_COLS = [\n",
        "    'config.w_init',\n",
        "    'config.activation',\n",
        "    'config.learning_rate',\n",
        "    'config.init_std',\n",
        "    'config.l2reg',\n",
        "    'config.train_fraction',\n",
        "    'config.dropout']\n",
        "CATEGORICAL_CONFIG_PARAMS = ['config.w_init', 'config.activation']\n",
        "CATEGORICAL_CONFIG_PARAMS_PREFIX = ['winit', 'act']\n",
        "DATAFRAME_METRIC_COLS = [\n",
        "    'test_accuracy',\n",
        "    'test_loss',\n",
        "    'train_accuracy',\n",
        "    'train_loss']\n",
        "TRAIN_SIZE = 15000\n",
        "\n",
        "# TODO: modify the following lines\n",
        "CONFIGS_PATH_BASE = 'path_to_the_file_with_best_configs'\n",
        "MNIST_OUTDIR = \"path_to_files_with_mnist_collection\"\n",
        "FMNIST_OUTDIR = 'path_to_files_with_fmnist_collection'\n",
        "CIFAR_OUTDIR = 'path_to_files_with_cifar10gs_collection'\n",
        "SVHN_OUTDIR = 'path_to_files_with_svhngs_collection'\n",
        "\n",
        "def filter_checkpoints(weights, dataframe,\n",
        "                       target='test_accuracy',\n",
        "                       stage='final', binarize=True):\n",
        "  \"\"\"Take one checkpoint per run and do some pre-processing.\n",
        "\n",
        "  Args:\n",
        "    weights: numpy array of shape (num_runs, num_weights)\n",
        "    dataframe: pandas DataFrame which has num_runs rows. First 4 columns should\n",
        "      contain test_accuracy, test_loss, train_accuracy, train_loss respectively.\n",
        "    target: string, what to use as an output\n",
        "    stage: flag defining which checkpoint out of potentially many we will take\n",
        "      for the run.\n",
        "    binarize: Do we want to binarize the categorical hyperparams?\n",
        "\n",
        "  Returns:\n",
        "    tuple (weights_new, metrics, hyperparams, ckpts), where\n",
        "      weights_new is a numpy array of shape (num_remaining_ckpts, num_weights),\n",
        "      metrics is a numpy array of shape (num_remaining_ckpts, num_metrics) with\n",
        "        num_metric being the length of DATAFRAME_METRIC_COLS,\n",
        "      hyperparams is a pandas DataFrame of num_remaining_ckpts rows and columns\n",
        "        listed in DATAFRAME_CONFIG_COLS.\n",
        "      ckpts is an instance of pandas Index, keeping filenames of the checkpoints\n",
        "    All the num_remaining_ckpts rows correspond to one checkpoint out of each\n",
        "    run we had.\n",
        "  \"\"\"\n",
        "\n",
        "  assert target in DATAFRAME_METRIC_COLS, 'unknown target'\n",
        "  ids_to_take = []\n",
        "  # Keep in mind that the rows of the DataFrame were sorted according to ckpt\n",
        "  # Fetch the unit id corresponding to the ckpt of the first row\n",
        "  current_uid = dataframe.axes[0][0].split('/')[-2]  # get the unit id\n",
        "  steps = []\n",
        "  for i in range(len(dataframe.axes[0])):\n",
        "    # Fetch the new unit id\n",
        "    ckpt = dataframe.axes[0][i]\n",
        "    parts = ckpt.split('/')\n",
        "    if parts[-2] == current_uid:\n",
        "      steps.append(int(parts[-1].split('-')[-1]))\n",
        "    else:\n",
        "      # We need to process the previous unit\n",
        "      # and choose which ckpt to take\n",
        "      steps_sort = sorted(steps)\n",
        "      target_step = -1\n",
        "      if stage == 'final':\n",
        "        target_step = steps_sort[-1]\n",
        "      elif stage == 'early':\n",
        "        target_step = steps_sort[0]\n",
        "      else:  # middle\n",
        "        target_step = steps_sort[int(len(steps) / 2)]\n",
        "      offset = [j for (j, el) in enumerate(steps) if el == target_step][0]\n",
        "      # Take the DataFrame row with the corresponding row id\n",
        "      ids_to_take.append(i - len(steps) + offset)\n",
        "      current_uid = parts[-2]\n",
        "      steps = [int(parts[-1].split('-')[-1])]\n",
        "\n",
        "  # Fetch the hyperparameters of the corresponding checkpoints\n",
        "  hyperparams = dataframe[DATAFRAME_CONFIG_COLS]\n",
        "  hyperparams = hyperparams.iloc[ids_to_take]\n",
        "  if binarize:\n",
        "    # Binarize categorical features\n",
        "    hyperparams = pd.get_dummies(\n",
        "        hyperparams,\n",
        "        columns=CATEGORICAL_CONFIG_PARAMS,\n",
        "        prefix=CATEGORICAL_CONFIG_PARAMS_PREFIX)\n",
        "  else:\n",
        "    # Make the categorical features have pandas type \"category\"\n",
        "    # Then LGBM can use those as categorical\n",
        "    hyperparams.is_copy = False\n",
        "    for col in CATEGORICAL_CONFIG_PARAMS:\n",
        "      hyperparams[col] = hyperparams[col].astype('category')\n",
        "\n",
        "  # Fetch the file paths of the corresponding checkpoints\n",
        "  ckpts = dataframe.axes[0][ids_to_take]\n",
        "\n",
        "  return (weights[ids_to_take, :],\n",
        "          dataframe[DATAFRAME_METRIC_COLS].values[ids_to_take, :].astype(\n",
        "              np.float32),\n",
        "          hyperparams,\n",
        "          ckpts)\n",
        "\n",
        "def build_fcn(n_layers, n_hidden, n_outputs, dropout_rate, activation,\n",
        "              w_regularizer, w_init, b_init, last_activation='softmax'):\n",
        "  \"\"\"Fully connected deep neural network.\"\"\"\n",
        "  model = keras.Sequential()\n",
        "  model.add(keras.layers.Flatten())\n",
        "  for _ in range(n_layers):\n",
        "    model.add(\n",
        "        keras.layers.Dense(\n",
        "            n_hidden,\n",
        "            activation=activation,\n",
        "            kernel_regularizer=w_regularizer,\n",
        "            kernel_initializer=w_init,\n",
        "            bias_initializer=b_init))\n",
        "    if dropout_rate \u003e 0.0:\n",
        "      model.add(keras.layers.Dropout(dropout_rate))\n",
        "  if n_layers \u003e 0:\n",
        "    model.add(keras.layers.Dense(n_outputs, activation=last_activation))\n",
        "  else:\n",
        "    model.add(keras.layers.Dense(\n",
        "        n_outputs,\n",
        "        activation='sigmoid',\n",
        "        kernel_regularizer=w_regularizer,\n",
        "        kernel_initializer=w_init,\n",
        "        bias_initializer=b_init))\n",
        "  return model\n",
        "\n",
        "def extract_summary_features(w, qts=(0, 25, 50, 75, 100)):\n",
        "  \"\"\"Extract various statistics from the flat vector w.\"\"\"\n",
        "  features = np.percentile(w, qts)\n",
        "  features = np.append(features, [np.std(w), np.mean(w)])\n",
        "  return features\n",
        "\n",
        "\n",
        "def extract_per_layer_features(w, qts=None, layers=(0, 1, 2, 3)):\n",
        "  \"\"\"Extract per-layer statistics from the weight vector and concatenate.\"\"\"\n",
        "  # Indices of the location of biases/kernels in the flattened vector\n",
        "  all_boundaries = {\n",
        "      0: [(0, 16), (16, 160)], \n",
        "      1: [(160, 176), (176, 2480)], \n",
        "      2: [(2480, 2496), (2496, 4800)], \n",
        "      3: [(4800, 4810), (4810, 4970)]}\n",
        "  boundaries = []\n",
        "  for layer in layers:\n",
        "    boundaries += all_boundaries[layer]\n",
        "  \n",
        "  if not qts:\n",
        "    features = [extract_summary_features(w[a:b]) for (a, b) in boundaries]\n",
        "  else:\n",
        "    features = [extract_summary_features(w[a:b], qts) for (a, b) in boundaries]\n",
        "  all_features = np.concatenate(features)\n",
        "  return all_features\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SBM6xNSjz8Bs"
      },
      "source": [
        "# 1. Loading the Small CNN Zoo dataset\n",
        "\n",
        "The following code loads the dataset (trained weights from *.npy files and all the relevant metrics, including accuracy, from *.csv files). "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mp-5POSc0ap2"
      },
      "outputs": [],
      "source": [
        "all_dirs = [MNIST_OUTDIR, FMNIST_OUTDIR, CIFAR_OUTDIR, SVHN_OUTDIR]\n",
        "weights = {'mnist': None,\n",
        "            'fashion_mnist': None,\n",
        "            'cifar10': None,\n",
        "            'svhn_cropped': None}\n",
        "metrics = {'mnist': None,\n",
        "            'fashion_mnist': None,\n",
        "            'cifar10': None,\n",
        "            'svhn_cropped': None}\n",
        "for (dirname, dataname) in zip(\n",
        "    all_dirs, ['mnist', 'fashion_mnist', 'cifar10', 'svhn_cropped']):\n",
        "  print('Loading %s' % dataname)\n",
        "  with gfile.GFile(os.path.join(dirname, \"all_weights.npy\"), \"rb\") as f:\n",
        "    # Weights of the trained models\n",
        "    weights[dataname] = np.load(f)\n",
        "  with gfile.GFile(os.path.join(dirname, \"all_metrics.csv\")) as f:\n",
        "    # pandas DataFrame with metrics\n",
        "    metrics[dataname] = pd.read_csv(f, index_col=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FNqWMZcx1y5m"
      },
      "source": [
        "Next it filters the dataset by keeping only checkpoints corresponding to 18 epochs and discarding runs that resulted in numerical instabilities. Finally, it performs the train / test splits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1qL5_-FZ11gm"
      },
      "outputs": [],
      "source": [
        "weights_train = {}\n",
        "weights_test = {}\n",
        "configs_train = {}\n",
        "configs_test = {}\n",
        "outputs_train = {}\n",
        "outputs_test = {}\n",
        "\n",
        "for dataset in ['mnist', 'fashion_mnist', 'cifar10', 'svhn_cropped']:\n",
        "  # Take one checkpoint per each run\n",
        "  # If using GBM as predictor, set binarize=False\n",
        "  weights_flt, metrics_flt, configs_flt, ckpts = filter_checkpoints(\n",
        "      weights[dataset], metrics[dataset], binarize=True)\n",
        "\n",
        "  # Filter out DNNs with NaNs and Inf in the weights\n",
        "  idx_valid = (np.isfinite(weights_flt).mean(1) == 1.0)\n",
        "  inputs = np.asarray(weights_flt[idx_valid], dtype=np.float32)\n",
        "  outputs = np.asarray(metrics_flt[idx_valid], dtype=np.float32)\n",
        "  configs = configs_flt.iloc[idx_valid]\n",
        "  ckpts = ckpts[idx_valid]\n",
        "\n",
        "  # Shuffle and split the data\n",
        "  random_idx = list(range(inputs.shape[0]))\n",
        "  np.random.shuffle(random_idx)\n",
        "  weights_train[dataset], weights_test[dataset] = (\n",
        "      inputs[random_idx[:TRAIN_SIZE]], inputs[random_idx[TRAIN_SIZE:]])\n",
        "  outputs_train[dataset], outputs_test[dataset] = (\n",
        "      1. * outputs[random_idx[:TRAIN_SIZE]],\n",
        "      1. * outputs[random_idx[TRAIN_SIZE:]])\n",
        "  configs_train[dataset], configs_test[dataset] = (\n",
        "      configs.iloc[random_idx[:TRAIN_SIZE]], \n",
        "      configs.iloc[random_idx[TRAIN_SIZE:]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "K7cpDNyB2tCc"
      },
      "source": [
        "# 2. Figure 2 of the paper\n",
        "\n",
        "Next we plot distribution of CNNs from 4 collections in Small CNN Zoo according to their train / test accuracy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VnToqYeT25pb"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize = (16, 8))\n",
        "pic_id = 0\n",
        "\n",
        "for dataset in ['mnist', 'fashion_mnist', 'cifar10', 'svhn_cropped']:\n",
        "  pic_id += 1\n",
        "  sp = plt.subplot(2, 4, pic_id)\n",
        "\n",
        "  outputs = outputs_train[dataset]\n",
        "\n",
        "  if dataset == 'mnist':\n",
        "    plt.title('MNIST', fontsize=24)\n",
        "  if dataset == 'fashion_mnist':\n",
        "    plt.title('Fashion MNIST', fontsize=24)\n",
        "  if dataset == 'cifar10':\n",
        "    plt.title('CIFAR10-GS', fontsize=24)\n",
        "  if dataset == 'svhn_cropped':\n",
        "    plt.title('SVHN-GS', fontsize=24)\n",
        "\n",
        "  # 1. test accuracy hist plots\n",
        "  sns.distplot(np.array(outputs[:, 0]), bins=15, kde=False, color='green')\n",
        "  plt.xlim((0.0, 1.0))\n",
        "  sp.axes.get_xaxis().set_ticklabels([])\n",
        "  sp.axes.get_yaxis().set_ticklabels([])\n",
        "  pic_id += 4\n",
        "  sp = plt.subplot(2, 4, pic_id)\n",
        "\n",
        "  # 2. test / train accuracy scatter plots\n",
        "  NUM_POINTS = 1000\n",
        "  random_idx = range(len(outputs))\n",
        "  np.random.shuffle(random_idx)\n",
        "  plt.plot([0.0, 1.0], [0.0, 1.0], 'r--')\n",
        "  sns.scatterplot(np.array(outputs[random_idx[:NUM_POINTS], 0]),  # test acc\n",
        "                  np.array(outputs[random_idx[:NUM_POINTS], 2]),  # train acc\n",
        "                  s=30\n",
        "                  )\n",
        "  if pic_id == 5:\n",
        "    plt.ylabel('Train accuracy', fontsize=22)\n",
        "    sp.axes.get_yaxis().set_ticklabels([0.0, 0.2, .4, .6, .8, 1.])\n",
        "  else:\n",
        "    sp.axes.get_yaxis().set_ticklabels([])\n",
        "  plt.xlim((0.0, 1.0))\n",
        "  plt.ylim((0.0, 1.0))\n",
        "  sp.axes.get_xaxis().set_ticks([0.0, 0.2, .4, .6, .8, 1.])\n",
        "  sp.axes.tick_params(axis='both', labelsize=18)\n",
        "  plt.xlabel('Test accuracy', fontsize=22)\n",
        "\n",
        "  pic_id -= 4\n",
        "\n",
        "plt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "fxtGdIK55t9B"
      },
      "source": [
        "# 3. Examples of training Logit-Linear / GBM / DNN predictors\n",
        "\n",
        "Next we train 3 models on all 4 CNN collections with the best hyperparameter configurations we found during our studies (documented in Table 2 and Section 4 of the paper).\n",
        "\n",
        "First, we load the best hyperparameter configurations we found.\n",
        "The file best_configs.json contains a list. \n",
        "Each entry of that list corresponds to the single hyperparameter configuration. \n",
        "It consists of: \n",
        "\n",
        "  (1) name of the CNN collection (mnist/fashion mnist/cifar10/svhn) \n",
        "  \n",
        "  (2) predictor type (linear/dnn/lgbm)\n",
        "  \n",
        "  (3) type of inputs, (refer to Table 2)\n",
        "  \n",
        "  (4) value of MSE you will get training with these settings, \n",
        "  \n",
        "  (5) dictionary of \"parameter name\"-\u003e \"parameter value\" for the given type of predictor."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "B7oCx5rr6y4D"
      },
      "outputs": [],
      "source": [
        "with gfile.GFile(os.path.join(CONFIGS_PATH_BASE, 'best_configs.json'), 'r') as file:\n",
        "  best_configs = json.load(file)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nQsP1aA5UhqT"
      },
      "source": [
        "# 3.1 Training GBM predictors\n",
        "\n",
        "GBM code below requires the lightgbm package.\n",
        "\n",
        "This is an example of training GBM on CIFAR10-GS CNN collection using per-layer weights statistics as inputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t4KzPiTAXWuo"
      },
      "outputs": [],
      "source": [
        "# Take the best config we found\n",
        "config = [el[-1] for el in best_configs if\n",
        "          el[0] == 'cifar10' and\n",
        "          el[1] == 'lgbm' and\n",
        "          el[2] == 'wstats-perlayer'][0]\n",
        "\n",
        "# Pre-process the weights\n",
        "train_x = np.apply_along_axis(\n",
        "    extract_per_layer_features, 1,\n",
        "    weights_train['cifar10'],\n",
        "    qts=None,\n",
        "    layers=(0, 1, 2, 3))\n",
        "test_x = np.apply_along_axis(\n",
        "    extract_per_layer_features, 1,\n",
        "    weights_test['cifar10'], \n",
        "    qts=None, \n",
        "    layers=(0, 1, 2, 3))\n",
        "# Get the target values\n",
        "train_y, test_y = outputs_train['cifar10'][:, 0], outputs_test['cifar10'][:, 0]\n",
        "\n",
        "# Define the GBM model\n",
        "lgbm_model = lgb.LGBMRegressor(\n",
        "    num_leaves=config['num_leaves'],\n",
        "    max_depth=config['max_depth'],\n",
        "    learning_rate=config['learning_rate'],\n",
        "    max_bin=int(config['max_bin']),\n",
        "    min_child_weight=config['min_child_weight'],\n",
        "    reg_lambda=config['reg_lambda'],\n",
        "    reg_alpha=config['reg_alpha'],\n",
        "    subsample=config['subsample'],\n",
        "    subsample_freq=1,  # it means always subsample\n",
        "    colsample_bytree=config['colsample_bytree'],\n",
        "    n_estimators=2000,\n",
        "    first_metric_only=True\n",
        ")\n",
        "\n",
        "# Train the GBM model;\n",
        "# Early stopping will be based on rmse of test set\n",
        "eval_metric = ['rmse', 'l1']\n",
        "eval_set = [(test_x, test_y)]\n",
        "lgbm_model.fit(train_x, train_y, verbose=100,\n",
        "               early_stopping_rounds=500,\n",
        "               eval_metric=eval_metric,\n",
        "               eval_set=eval_set,\n",
        "               eval_names=['test'])\n",
        "\n",
        "# Evaluate the GBM model\n",
        "assert hasattr(lgbm_model, 'best_iteration_')\n",
        "# Choose the step which had the best rmse on the test set\n",
        "best_iter = lgbm_model.best_iteration_ - 1\n",
        "lgbm_history = lgbm_model.evals_result_\n",
        "mse = lgbm_history['test']['rmse'][best_iter] ** 2.\n",
        "mad = lgbm_history['test']['l1'][best_iter]\n",
        "var = np.mean((test_y - np.mean(test_y)) ** 2.)\n",
        "r2  = 1. - mse / var\n",
        "print('Test MSE = ', mse)\n",
        "print('Test MAD = ', mad)\n",
        "print('Test R2 = ', r2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1Sf5cFosZcmk"
      },
      "source": [
        "# 3.2 Training DNN predictors\n",
        "\n",
        "This is an example of training DNN on MNIST CNN collection using all weights as inputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cVsPbhQYZodD"
      },
      "outputs": [],
      "source": [
        "# Take the best config we found\n",
        "config = [el[-1] for el in best_configs if\n",
        "          el[0] == 'mnist' and\n",
        "          el[1] == 'dnn' and\n",
        "          el[2] == 'weights'][0]\n",
        "\n",
        "train_x, test_x = weights_train['cifar10'], weights_test['cifar10']\n",
        "train_y, test_y = outputs_train['cifar10'][:, 0], outputs_test['cifar10'][:, 0]\n",
        "\n",
        "# Get the optimizer, initializers, and regularizers\n",
        "optimizer = keras.optimizers.get(config['optimizer_name'])\n",
        "optimizer.learning_rate = config['learning_rate']\n",
        "w_init = keras.initializers.get(config['w_init_name'])\n",
        "if config['w_init_name'].lower() in ['truncatednormal', 'randomnormal']:\n",
        "  w_init.stddev = config['init_stddev']\n",
        "b_init = keras.initializers.get('zeros')\n",
        "w_reg = (keras.regularizers.l2(config['l2_penalty']) \n",
        "         if config['l2_penalty'] \u003e 0 else None)\n",
        "\n",
        "# Get the fully connected DNN architecture\n",
        "dnn_model = build_fcn(int(config['n_layers']),\n",
        "                      int(config['n_hiddens']),\n",
        "                      1,  # number of outputs\n",
        "                      config['dropout_rate'],\n",
        "                      'relu',\n",
        "                      w_reg, w_init, b_init,\n",
        "                      'sigmoid')  # Last activation\n",
        "dnn_model.compile(\n",
        "    optimizer=optimizer,\n",
        "    loss='mean_squared_error',\n",
        "    metrics=['mse', 'mae'])\n",
        "\n",
        "# Train the model\n",
        "dnn_model.fit(\n",
        "    train_x, train_y,\n",
        "    batch_size=int(config['batch_size']),\n",
        "    epochs=300,\n",
        "    validation_data=(test_x, test_y),\n",
        "    verbose=1,\n",
        "    callbacks=[keras.callbacks.EarlyStopping(\n",
        "        monitor='val_loss',\n",
        "        min_delta=0,\n",
        "        patience=10,\n",
        "        verbose=0,\n",
        "        mode='auto',\n",
        "        baseline=None,\n",
        "        restore_best_weights=False)]\n",
        "    )\n",
        "\n",
        "# Evaluate the model\n",
        "eval_train = dnn_model.evaluate(train_x, train_y, batch_size=128, verbose=0)\n",
        "eval_test = dnn_model.evaluate(test_x, test_y, batch_size=128, verbose=0)\n",
        "assert dnn_model.metrics_names[1] == 'mean_squared_error'\n",
        "assert dnn_model.metrics_names[2] == 'mean_absolute_error'\n",
        "mse = eval_test[1]\n",
        "var = np.mean((test_y - np.mean(test_y)) ** 2.)\n",
        "r2  = 1. - mse / var\n",
        "print('Test MSE = ', mse)\n",
        "print('Test MAD = ', eval_test[2])\n",
        "print('Test R2 = ', r2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DF3N5jZ9JQMs"
      },
      "source": [
        "# 3.3 Train Logit-Linear predictors\n",
        "\n",
        "This is an example of training Logit-Linear model on CIFAR10 CNN collection using hyperparameters as inputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_S_183RnJUZu"
      },
      "outputs": [],
      "source": [
        "# Take the best config we found\n",
        "config = [el[-1] for el in best_configs if\n",
        "          el[0] == 'cifar10' and\n",
        "          el[1] == 'linear' and\n",
        "          el[2] == 'hyper'][0]\n",
        "\n",
        "# Turn DataFrames to numpy arrays. \n",
        "# Since we used \"binarize=True\" when calling filter_checkpoints all the\n",
        "# categorical columns were binarized.\n",
        "train_x = configs_train['cifar10'].values.astype(np.float32)\n",
        "test_x = configs_test['cifar10'].values.astype(np.float32)\n",
        "train_y, test_y = outputs_train['cifar10'][:, 0], outputs_test['cifar10'][:, 0]\n",
        "\n",
        "# Get the optimizer, initializers, and regularizers\n",
        "optimizer = keras.optimizers.get(config['optimizer_name'])\n",
        "optimizer.learning_rate = config['learning_rate']\n",
        "w_init = keras.initializers.get(config['w_init_name'])\n",
        "if config['w_init_name'].lower() in ['truncatednormal', 'randomnormal']:\n",
        "  w_init.stddev = config['init_stddev']\n",
        "b_init = keras.initializers.get('zeros')\n",
        "w_reg = (keras.regularizers.l2(config['l2_penalty']) \n",
        "         if config['l2_penalty'] \u003e 0 else None)\n",
        "\n",
        "# Get the linear architecture (DNN with 0 layers)\n",
        "dnn_model = build_fcn(int(config['n_layers']),\n",
        "                      int(config['n_hiddens']),\n",
        "                      1,  # number of outputs\n",
        "                      None,  # Dropout is not used\n",
        "                      'relu',\n",
        "                      w_reg, w_init, b_init,\n",
        "                      'sigmoid')  # Last activation\n",
        "dnn_model.compile(\n",
        "    optimizer=optimizer,\n",
        "    loss='mean_squared_error',\n",
        "    metrics=['mse', 'mae'])\n",
        "\n",
        "# Train the model\n",
        "dnn_model.fit(\n",
        "    train_x, train_y,\n",
        "    batch_size=int(config['batch_size']),\n",
        "    epochs=300,\n",
        "    validation_data=(test_x, test_y),\n",
        "    verbose=1,\n",
        "    callbacks=[keras.callbacks.EarlyStopping(\n",
        "        monitor='val_loss',\n",
        "        min_delta=0,\n",
        "        patience=10,\n",
        "        verbose=0,\n",
        "        mode='auto',\n",
        "        baseline=None,\n",
        "        restore_best_weights=False)]\n",
        "    )\n",
        "\n",
        "# Evaluate the model\n",
        "eval_train = dnn_model.evaluate(train_x, train_y, batch_size=128, verbose=0)\n",
        "eval_test = dnn_model.evaluate(test_x, test_y, batch_size=128, verbose=0)\n",
        "assert dnn_model.metrics_names[1] == 'mean_squared_error'\n",
        "assert dnn_model.metrics_names[2] == 'mean_absolute_error'\n",
        "mse = eval_test[1]\n",
        "var = np.mean((test_y - np.mean(test_y)) ** 2.)\n",
        "r2  = 1. - mse / var\n",
        "print('Test MSE = ', mse)\n",
        "print('Test MAD = ', eval_test[2])\n",
        "print('Test R2 = ', r2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "i97PjpsxWQWS"
      },
      "source": [
        "# 4. Figure 4: Transfer across datasets\n",
        "\n",
        "Train GBM predictor using statistics of all layers as inputs on all 4 CNN collections. Then evaluate them on each of the 4 CNN collections (without fine-tuning). Store all results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "xRFiVulhWeQ9"
      },
      "outputs": [],
      "source": [
        "transfer_results = {}\n",
        "\n",
        "for dataset in ['mnist', 'fashion_mnist', 'cifar10', 'svhn_cropped']:\n",
        "  print('Training on %s' % dataset)\n",
        "  transfer_results[dataset] = {}\n",
        "  \n",
        "  train_x = weights_train[dataset]\n",
        "  test_x = weights_test[dataset]\n",
        "  train_y = outputs_train[dataset][:, 0]\n",
        "  test_y = outputs_test[dataset][:, 0]\n",
        "\n",
        "  # Pre-process the weights by taking the statistics across layers\n",
        "  train_x = np.apply_along_axis(\n",
        "      extract_per_layer_features, 1, \n",
        "      train_x, qts=None, layers=(0, 1, 2, 3))\n",
        "  test_x = np.apply_along_axis(\n",
        "      extract_per_layer_features, 1,\n",
        "      test_x, qts=None, layers=(0, 1, 2, 3))\n",
        "\n",
        "  # Take the best config we found\n",
        "  config = [el[-1] for el in best_configs if\n",
        "            el[0] == dataset and\n",
        "            el[1] == 'lgbm' and\n",
        "            el[2] == 'wstats-perlayer'][0]\n",
        "\n",
        "  lgbm_model = lgb.LGBMRegressor(\n",
        "      num_leaves=config['num_leaves'],\n",
        "      max_depth=config['max_depth'], \n",
        "      learning_rate=config['learning_rate'], \n",
        "      max_bin=int(config['max_bin']),\n",
        "      min_child_weight=config['min_child_weight'],\n",
        "      reg_lambda=config['reg_lambda'],\n",
        "      reg_alpha=config['reg_alpha'],\n",
        "      subsample=config['subsample'],\n",
        "      subsample_freq=1,  # Always subsample\n",
        "      colsample_bytree=config['colsample_bytree'],\n",
        "      n_estimators=4000,\n",
        "      first_metric_only=True,\n",
        "      )\n",
        "  \n",
        "  # Train the GBM model\n",
        "  lgbm_model.fit(\n",
        "      train_x,\n",
        "      train_y,\n",
        "      verbose=100,\n",
        "      # verbose=False,\n",
        "      early_stopping_rounds=500,\n",
        "      eval_metric=['rmse', 'l1'],\n",
        "      eval_set=[(test_x, test_y)],\n",
        "      eval_names=['test'])\n",
        "  \n",
        "  # Evaluate on all 4 CNN collections\n",
        "  for transfer_to in ['mnist', 'fashion_mnist', 'cifar10', 'svhn_cropped']:\n",
        "    print('Evaluating on %s' % transfer_to)\n",
        "    # Take the test split of the dataset\n",
        "    transfer_x = weights_test[transfer_to]\n",
        "    transfer_x = np.apply_along_axis(\n",
        "        extract_per_layer_features, 1,\n",
        "        transfer_x, qts=None, layers=(0, 1, 2, 3))\n",
        "    y_hat = lgbm_model.predict(transfer_x)\n",
        "    transfer_results[dataset][transfer_to] = y_hat"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VvkJS4CKYDj_"
      },
      "source": [
        "And plot everything"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "U9J8nA4BYF4P"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize = (15, 15))\n",
        "pic_id = 0\n",
        "for dataset in ['mnist', 'fashion_mnist', 'cifar10', 'svhn_cropped']:\n",
        "  for transfer_to in ['mnist', 'fashion_mnist', 'cifar10', 'svhn_cropped']:\n",
        "    pic_id += 1\n",
        "    sp = plt.subplot(4, 4, pic_id)\n",
        "    # Take true labels\n",
        "    y_true = outputs_test[transfer_to][:, 0]\n",
        "    # Take the predictions of the model\n",
        "    y_hat = transfer_results[dataset][transfer_to]\n",
        "    plt.plot([0.01, .99], [0.01, .99], 'r--', linewidth=2)\n",
        "    sns.scatterplot(y_true, y_hat)\n",
        "    # Compute the Kendall's tau coefficient\n",
        "    tau = stats.kendalltau(y_true, y_hat)[0]\n",
        "    plt.text(0.05, 0.9, r\"$\\tau=%.3f$\" % tau, fontsize=25)\n",
        "    plt.xlim((0.0, 1.0))\n",
        "    plt.ylim((0.0, 1.0))\n",
        "\n",
        "    if pic_id % 4 != 1:\n",
        "      sp.axes.get_yaxis().set_ticklabels([])\n",
        "    else:\n",
        "      plt.ylabel('Predictions', fontsize=22)\n",
        "      sp.axes.tick_params(axis='both', labelsize=15)\n",
        "\n",
        "    if pic_id \u003c 13:\n",
        "      sp.axes.get_xaxis().set_ticklabels([])\n",
        "    else:\n",
        "      plt.xlabel('Test accuracy', fontsize=22)\n",
        "      sp.axes.tick_params(axis='both', labelsize=15)\n",
        "\n",
        "    if pic_id == 1:\n",
        "      plt.title('MNIST', fontsize=22)\n",
        "    if pic_id == 2:\n",
        "      plt.title('Fashion-MNIST', fontsize=22)\n",
        "    if pic_id == 3:\n",
        "      plt.title('CIFAR10-GS', fontsize=22)\n",
        "    if pic_id == 4:\n",
        "      plt.title('SVHN-GS', fontsize=22)\n",
        "\n",
        "plt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Iahn92bHY8kQ"
      },
      "source": [
        "# 5. Figure 3: various 2d plots based on subsets of weights statistics\n",
        "\n",
        "Take weight statistics for the CIFAR10 CNN collection. Plot various 2d plots"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nBXxv0-2ZfZA"
      },
      "outputs": [],
      "source": [
        "# Take the per-layer weights stats for the train split of CIFAR10-GS collection\n",
        "per_layer_stats = np.apply_along_axis(\n",
        "    extract_per_layer_features, 1,\n",
        "    weights_train['cifar10'])\n",
        "train_test_accuracy = outputs_train['cifar10'][:, 0]\n",
        "# Positions of various stats\n",
        "b0min = 0  # min of the first layer\n",
        "b0max = 4  # max of the first layer\n",
        "bnmin = 6*7 + 0  # min of the last layer\n",
        "bnmax = 6*7 + 4  # max of the last layer\n",
        "x = per_layer_stats[:,b0max] - per_layer_stats[:,b0min]\n",
        "y = per_layer_stats[:,bnmax] - per_layer_stats[:,bnmin]\n",
        "\n",
        "plt.figure(figsize=(10,8))\n",
        "plt.scatter(x, y, s=15,\n",
        "            c=train_test_accuracy,\n",
        "            cmap=\"jet\",\n",
        "            vmin=0.1,\n",
        "            vmax=0.54,\n",
        "            linewidths=0)\n",
        "plt.yscale(\"log\")\n",
        "plt.xscale(\"log\")\n",
        "plt.ylim(0.1, 10)\n",
        "plt.xlim(0.1, 10)\n",
        "plt.xlabel(\"Bias range, first layer\", fontsize=22)\n",
        "plt.ylabel(\"Bias range, final layer\", fontsize=22)\n",
        "cbar = plt.colorbar()\n",
        "cbar.ax.tick_params(labelsize=18) \n",
        "plt.tight_layout()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/brain/python/client:colab_notebook",
        "kind": "private"
      },
      "name": "dnn-predict-accuracy.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1AV92_u26P4KyTmopFOKgROWg3GostHAA",
          "timestamp": 1581610553676
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
